import os
import cv2
import sys
from .transformations import resize_transform_basic
import numpy as np
import pandas as pd
import cv2
from pathlib import Path
import torch
from torch.utils.data import DataLoader, Dataset, sampler
from utils.json import json_to_image


class T2_Seg_Dataset(Dataset):
    '''
    T2_Seg_Dataset
    '''
    def __init__(self, dataframe, transform=None, mask_label='box'):
        self.df = dataframe
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform
        self.mask_label = mask_label

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        maskPath = self.df["label"].iloc[idx]
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except IOError:
            print(imgPath)
            raise ValueError("The image path is incorrect, please check the image path.")
    
        if maskPath == str(0):
            mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
        else:
            if Path(maskPath).suffix == '.png':
                mask = cv2.imread(maskPath, 0) # only channle 2 has values
                mask = mask.astype(np.float32) / 255
            elif Path(maskPath).suffix == '.json':
                try:
                    mask = json_to_image(image.shape, maskPath, mask_label=self.mask_label)
                except:
                    print(imgPath)
                    raise ValueError("The image file %s has something wrong."%imgPath)
        
        h, w = image.shape[:2]
        if (h==1020 and w==1360):
            image = image[:962, :]
            mask = mask[:962, :]
        elif (h==864 and w==1176):
            new_w = int(w*394/359)
            image = cv2.resize(image, (new_w, int(h*394/359)), interpolation=cv2.INTER_LINEAR)
            mask = cv2.resize(mask, (new_w, int(h*394/359)), interpolation=cv2.INTER_NEAREST)
            mask[mask>0] = 1
            image = image[:919, :]
            mask = mask[:919, :]
            image = cv2.copyMakeBorder(image, 0, 0, 0, 1305-new_w, cv2.BORDER_CONSTANT, (0, 0, 0))
            mask = cv2.copyMakeBorder(mask, 0, 0, 0, 1305-new_w, cv2.BORDER_CONSTANT, 0)

        aug = self.seg_transforms(image=image, mask=mask)
        aug_image = aug['image']
        aug_mask = aug['mask']
        aug_mask = aug_mask.unsqueeze(0).float()
        return aug_image, aug_mask


class T2_Seg_Test_Dataset(Dataset):
    '''
    T2_Seg_Dataset
    '''
    def __init__(self, dataframe, transform=None, mask_label='box'):
        self.df = dataframe
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform
        self.mask_label = mask_label

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        label = self.df["label"].iloc[idx]
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except IOError:
            print(imgPath)
            raise ValueError("The image path is incorrect, please check the image path.")
        
        h, w = image.shape[:2]
        if (h==1020 and w==1360):
            image = image[:962, :]
        elif (h==864 and w==1176):
            new_w = int(w*394/359)
            image = cv2.resize(image, (new_w, int(h*394/359)), interpolation=cv2.INTER_LINEAR)
            image = image[:919, :]
            image = cv2.copyMakeBorder(image, 0, 0, 0, 1305-new_w, cv2.BORDER_CONSTANT, (0, 0, 0))

        aug = self.seg_transforms(image=image)
        aug_image = aug['image']
        # print('agu aug_image shape: ', aug_image.shape)
        # print('agu mask shape: ', aug_mask.shape)
        return aug_image, label, Path(imgPath).stem, imgPath

    
class T7_Mask_Dataset(Dataset):
    '''
    T7_Mask_Dataset
    '''
    def __init__(self, dataframe, transform=None):
        self.df = dataframe
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        label = self.df["label"].iloc[idx]
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except IOError:
            print(imgPath)
            raise IOError("The image path is incorrect, please check the image path.")

        aug = self.seg_transforms(image=image)

        if os.path.isfile(label):
            label = Path(label).parts[-2]
            
        return aug['image'], label, Path(imgPath).stem, imgPath

    
